﻿#!/usr/bin/env python
# -*- coding:utf-8 -*-
import time
import thread

from lib.mongo import mongo
from lib.sqlapi import sqlmapOpr
import json
import requests
import pymongo
__author__ = '云絮'

class scandispatcher:

    def __init__(self):
        self.db=mongo().db
        self.sqlapi=sqlmapOpr()
    '''
    开始扫描，从数据库中拉取5条任务，去扫描。任务中永远保持5条，没有会拉新数据
    '''
    def scanner(self):
        while True:
            count=self.db.taskmgr.find({'status':0}).count()
            #print count
            if count>5:
                continue
            try:
                task=self.db.tasklist.find_one({'isscan':0},{'isscan':0,'_id':0})
            except Exception,e:
                pass

            #print tasks[0]s
            if not task:
                continue
            conf=self.db.conf.find_one()
            agent=conf['agent']
            if agent:
                task['agent']=agent

            try:
                taskid=self.sqlapi.taskNew()
                self.sqlapi.taskStartByJson(taskid,task)
                url=task['url']
                print task
                #print '%s[%s] scaning...'%(url,taskid)
            except requests.exceptions.ConnectionError,e:
                exit('无法连接到sqlmapAPI')
            self.db.tasklist.update({'url':url},{'$set':{'isscan':1}},multi=True)
            self.db.taskmgr.insert({'taskid':taskid,'status':0})
            print 'insert task:'+taskid
            thread.start_new_thread(self.updateStatus,(taskid,url))
    #更新已经完成的任务，设置status为1
    def updateStatus(self,taskid,url):
        while True:
            status=self.sqlapi.getStatus(taskid)
            #如果还在跑注入，就等待10秒后在来检查
            if status=='running':
                time.sleep(10)
            #如果是已经完成，就设置为完成，更新状态为1，然后退出循环
            elif status=='terminated' or status=="not running":
                self.db.taskmgr.update({'taskid':taskid},{'$set':{'status':1}})
                #print '%s.scan Done'%taskid
                self.addScanResult(taskid,url)
                break
            #如果是其他提示，也退出循环，状态为2
            else:
                self.db.taskmgr.update({'taskid':taskid},{'$set':{'status':2}})
                break
    #通过扫描
    def addScanResult(self,taskid,url):
        #查询扫描结果，通过taskid
        result=json.loads(self.sqlapi.getScanResult(taskid))
        if result['data']:
            result={
                'url':'{url}'.format(url=url),
                'taskid':'{taskid}'.format(taskid=taskid),
                #'dbms':'{dbms}'.format(dbms=result['data'][0]["value"][0]["dbms"]),
                'method':'{method}'.format(method=result['data'][0]["value"][0]["place"]),
                'os':'{os}'.format(os=result['data'][0]["value"][0]['os']),
                'parameter':'{parameter}'.format(parameter=result['data'][0]["value"][0]['parameter']),
                'payload':'{payload}'.format(payload=result['data'][0]["value"][0]["data"]['1']['payload'])
            }
            #print result
            postData=''
            if result['method']=='POST':
                data=self.db.tasklist.find_one({'url':url})
                postData=data['data']
            result['postData']=postData
            print 'sql scan:',result
            #结果插入数据库
            self.db.resultlist.insert(result)

if __name__ == '__main__':
    scandispatcher().scanner()

